#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jul 19 01:01:42 2018
@project: 天池比赛-A股主板上市公司公告信息抽取
@group: MZH_314
@author: LHQ
"""
import re
import os
from collections import OrderedDict

import pandas as pd

from reportIE.preprocess import ZjcReport as Report
from reportIE.preprocess import make_converter
from reportIE.utils import load_reports, load_train_data
from reportIE.utils.qb import strQ2B


price_cvt = make_converter("价格")
date_cvt = make_converter("日期")
percent_cvt = make_converter("百分比")
quantity_cvt = make_converter("股数")


def df2item(df):
    items = df.to_dict("list", into=OrderedDict)
    for k, v in items.items():
        yield k, v
        
        
def strip(s):
    if isinstance(s, str):
        s = re.sub("\s", "", s)
    return s


def process_zjc_table(html, report_id):
    report = Report(html, report_id)
    origin_tables = report.tables
    for t_i, tb in enumerate(origin_tables):
        df_desc = tb.desc
        cols = []
        dftb = tb.to_df()
        if dftb.empty:
            continue
        for col_name in dftb.columns:
            col = dftb[col_name]
            if len(col.shape) > 1:
                continue
            col = col.map(strip)
            if re.search("(?:[(（]股[)）]|股数|股份)(?!时间)", col_name):
                col = col.map(quantity_cvt)
            if re.search("(?:比例|%|占.*?比)", col_name):
                col = col.map(percent_cvt)
            if re.search("(?:元|均价|价)", col_name):
                col = col.map(price_cvt)
            if re.search("(?:时间|期间|日期|期)(?!元|均价|价)", col_name):
                col = col.map(date_cvt)
            cols.append(col)
        if len(cols) < 1:
            continue
        df_new = pd.concat(cols, axis=1)
        for f_i, (k, v) in enumerate(df2item(df_new)):
            item = OrderedDict()
            item["report_id"] = report_id
            item["table_index"] = t_i
            item["table_desc"] = df_desc
            item["field_index"] = f_i
            item["field"] = k
            item["values"] = [strQ2B(str(ele)) for ele in v]
            yield item


def process_zjc_tables(html_dir):
    data = []
    for report_id, html in load_reports(html_dir):
        for item in process_zjc_table(html, report_id):
            data.append(item)
    df = pd.DataFrame(data)
    return df


if __name__ == "__main__":
    columns = ["公告id", "股东全称", "股东简称", "变动截止日期", "变动价格", "变动数量", "变动后持股数", "变动后持股比例"]
    train_path = os.path.abspath("../data/[new] FDDC_announcements_round1_train_result_20180616/zengjianchi.train")
    html_dir = os.path.abspath("../data/round1_train_20180518/增减持/html")

    save_path = os.path.abspath("../data/training_data/zengjianchi_table.csv")
    save_dir = os.path.dirname(save_path)
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
 
    df = process_zjc_tables(html_dir)

    # 标注训练集
    df_train = load_train_data(train_path, columns=columns)
    data = []
    for report_id, df_std in df_train.groupby('公告id'):
        shareholder_fullname = set(df_std["股东全称"].dropna().map(str).tolist())
        shareholder_shortname = set(df_std['股东简称'].dropna().map(str).tolist())
        change_date = set(df_std['变动截止日期'].dropna().map(str).tolist())
        change_price = set(df_std['变动价格'].dropna().map(str).tolist())
        change_quantity = set(df_std['变动数量'].dropna().map(str).tolist())
        final_quantity = set(df_std['变动后持股数'].dropna().map(int).map(str).tolist())
        final_percent = set(df_std['变动后持股比例'].dropna().map(str).tolist())
        
        report_id_str = str(report_id)
        df_origin = df.query("report_id == @report_id_str")
        if df_origin.empty:
            continue
        
        def label(row):
            values = row['values']
            vset = set([strQ2B(str(ele)) for ele in values if ele])
            row['label'] = '其他'
            if len(vset & shareholder_fullname) > 0:
                row['label'] = '股东全称'
            if len(vset & shareholder_shortname) > 0:
                row['label'] = '股东简称'
            if len(vset & change_date) > 0:
                row['label'] = '变动截止日期'
            if len(vset & change_price) > 0:
                row['label'] = '变动价格'
            if len(vset & change_quantity) > 0:
                row['label'] = '变动数量'
            if len(vset & final_quantity) > 0:
                row['label'] = '变动后持股数'
            if len(vset & final_percent) > 0:
                row['label'] = '变动后持股比例'
            return row
                
        df_label = df_origin.apply(label, axis=1)
        data.append(df_label)
#    
    df_new = pd.concat(data)
    df_new.to_csv(save_path, index=False)
#    
# 
